# 代码编写者： 李朔
# 开发时间： 2021/7/12 13:56
import torch
import matplotlib.pyplot as plt
#第一部分：数据集
x_data=torch.Tensor([[1.0],[2.0],[3.0]])
y_data=torch.Tensor([[2.0],[4.0],[6.0]])

#第二部分：model
class LinearModel(torch.nn.Module):
    def __init__(self):
        super(LinearModel,self).__init__()
        self.linear=torch.nn.Linear(1,1)  #创建Linear对象实例，参数为x,y的维数

    def forward(self,x):
        y_pred=self.linear(x)
        return y_pred
#将对象实例化
model=LinearModel()

#第三部分：损失函数以及优化器
criterion=torch.nn.MSELoss(size_average=False)
optimizer=torch.optim.SGD(model.parameters(),lr=0.01)

#第四部分：training cycle
loss_list=[]
for epoch in range(100):
    y_pred=model(x_data)
    loss=criterion(y_pred,y_data)
    loss_list.append(loss.item())
    print("epoch:{0}\t loss:{1}".format(epoch,loss.item()))

    optimizer.zero_grad()
    loss.backward()    #反馈函数
    optimizer.step()    #更新
print("w=",model.linear.weight.item())
print("b=",model.linear.bias.item())

#测试集
x_test=torch.Tensor([[4.0]])
y_test=model(x_test)
print("y_pred=",y_test.data)
#绘图
plt.plot(loss_list)
plt.xlabel("epoch")
plt.ylabel("loss")
plt.show()